#!/usr/bin/env python
# -*- coding:utf-8 -*-
import json
import os

from langchain_core.prompts import ChatPromptTemplate

from controller.pptGenController import PptGenController
from models import db
from models.ppt_base import PptBase
from models.ppt_gen import PptGen
from service.llmService import LLMService
from service.pptBaseService import PptBaseService
from service.pptChatDocumentService import PptChatDocumentService
from service.templateService import TemplateService
from utils.analysis_ppt import ppt_structure_parser
from utils.fetch_json import fetch_json
from utils.generate_id import GenerateID
from utils.loggings import loggings
from utils.ppt_element_enums import PPT_ELEMENT_ENUMS
from utils.ppt_generation import PPTGeneration
from utils.response_code import RET


class PptGenService(PptGenController):
    base_dir = os.path.dirname(os.path.dirname(__file__))
    ppt_save_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'uploads', 'generation')

    element_analysis_prompt = """你是一个PPT模板分析师，现在用户给你一个PPT模板中每一页的结构，包括了每一页的标题、内容、图片、表格等元素。
你要做的是基于用户的输入，分析这个PPT模板的结构，从而判断每一页适合呈现哪种内容，比如哪一页适合放标题，哪一页是目录，哪一页是章节的总目录等。
判断类型如下：
{ppt_elements_enum}

对于文字内容页模板、图片内容页模板以及表格页模板，你需要判断这一页能够放多少项内容，例如结构中如果有三项文本框等信息，你可以认为可以放3项内容。
以下是某一页的结构示例：
elements_info: [{{{{'type': 'TextBox', 'content': '单击此处输入你的项正文，文字是您思想的提炼，请尽量言简意赅的阐述观点，单击此处输入你的项正文'}}}},
{{{{'type': 'TextBox', 'content': '单击此处输入你的项正文，文字是您思想的提炼，请尽量言简意赅的阐述观点，单击此处输入你的项正文'}}}},
{{{{'type': 'TextBox', 'content': '单击此处输入你的项正文，文字是您思想的提炼，请尽量言简意赅的阐述观点，单击此处输入你的项正文'}}}},
{{{{'type': 'TextBox', 'content': '01'}}}},
{{{{'type': 'TextBox', 'content': '02'}}}},
{{{{'type': 'TextBox', 'content': '03'}}}},
{{{{'type': 'TextBox', 'content': '添加项标题'}}}},
{{{{'type': 'TextBox', 'content': '单击此处添加标题'}}}}]

这里内容项的数量是3，因为有'01,02,03'这三个固有的一些界面元素，不能修改，配合'单击此处输入你的项正文...'说明是一组内容项，那么这就是一个可以放3项内容的模板页。而不是8项内容。
标题和项标题等不纳入计数。你需要根据这个信息，判断每一页的模板类型，以及适用的文本项数量。


请输出一个json格式的结果，示例如下：
[{{{{
    "page": 1,
    "type": 1,
    "item": null,
}}}}, {{{{
    "page": 2,
    "type": 2,
    "item": null,
}}}}, {{{{
    "page": 3,
    "type": 3,
    "item": null,
}}}}, {{{{
    "page": 4,
    "type": 4,
    "item": 3,
}}}}, {{{{
    "page": 5,
    "type": 5,
    "item": 4,
}}}}]
"""

    plan_ppt_prompt = """你是一个PPT设计师，用户会给你一个大纲，以及PPT模板的信息，包括了每一页的模板类型以及适用的文本项数量。
你需要做的是基于用户的输入，规划这个PPT的内容，从而能够将内容填充到对应的模板中，生成一个完整的PPT文件，为了最终产生一个美观的PPT，你要合理地规划每一页模板，保证模板与内容能够适配。
当然，尽可能地使用不同的模板页，不过也允许重新使用某一页模板。你要注意，模板页不是所有的页都必须使用上，而是你要基于大纲合理地规划模板页的使用。
模板类型如下：
{ppt_elements_enum}

用户给出的PPT模板元素信息如下：
"page": 模板页
"type": 模板类型
"item": 有多少文本项，你要根据这个来规划内容，例如如果有三项文本框，你可以认为必须放3项内容

注意事项：
1. 用户给出的每一页的结构，包括了每一页的标题、内容、图片、表格等元素，其中有一些是不能修改的，比如说01、02、03等，这是固有的一些界面元素，你不能修改，而有一些你可以修改，你需要进行判断区分。
2. 你要注意合理地分配每一页模板，但是对于相同的内容，只需要选择某一个模板页，而不是每一类模板都使用一遍。不是所有的模板页都必须被使用，而是根据实际的内容来选择模板！
4. 不要重复阐述相同的内容，只需要使用一种模板并设计计划即可！例如阐述`XXX安全性`，不要在多个页面重复阐述安全性，只需要选择一个页面阐述即可！相同模板类型的页面不要重复使用！
5. 核心是根据内容选择模板类型与模板页，而不是对每个模板页都要适配内容！模板页允许重复使用。

现在你需要基于用户给出的PPT内容大纲，以及每一页的模板类型和适用的文本项数量，规划这个PPT的内容，给出PPT的模板使用规划，要满足注意事项中的要求。

以下是用户给出的PPT设计计划的格式说明：
page: 实际生成的界面页码
template_page: 模板页码，注意模板页按照之前的要求是可以重复使用的
template_type: 模板类型
plan: 计划方案

请输出一个json格式的结果，示例如下：
[{{{{
    "page": 1,
    "template_page": 1,
    ”template_type": 1,
    "plan": "需要填入标题内容"
}}}}, {{{{
    "page": 2,
    "template_page": 2,
    ”template_type": 2,
    "plan": "需要填入目录内容：引言\n第一章\n第二章\n第三章"
}}}}, {{{{
    "page": 3,
    "template_page": 3,
    ”template_type": 3,
    "plan": "需要填入子标题：引言；需要修改索引为0"
}}}}, {{{{
    "page": 4,
    "template_page": 4,
    ”template_type": 4,
    "plan": "需要填入引言的子标题内容，A，同时围绕主题A展开3项内容，满足顺序要求"
}}}}, {{{{
    "page": 5,
    "template_page": 15,
    ”template_type": 5,
    "content": "需要填入引言的子标题内容，B，同时围绕主题B展开3项内容，满足顺序要求。另外采用了图文模板，还需要插入三张图片"
}}}}, {{{{
    "page": 6,
    "template_page": 3,
    ”template_type": 3,
    "plan": "需要填入子标题：第一章"
}}}}, {{{{
    "page": 7,
    "template_page": 6,
    ”template_type": 4,
    "plan": "需要填入第一章的子标题内容，C，同时围绕主题C展开4项内容，按照并列关系"
}}}}]
    
"""

    single_ppt_generation = """你是一个PPT制作专家，你会基于当前页面的模板元素，以及用户给出的PPT设计计划来搜索并填充内容。以下是用户给出的信息：
```
当前页面的模板元素：
{ppt_element}

其中的元素信息包括：
text_boxes: 文本框数量
images: 图片数量
tables: 表格数量
elements_info: 实际的spire.presentation代码元素

最终要基于elements_info中的元素信息，替换对应的文本/图片/表格内容，你要知道每一个模板元素，同时你要区分哪些是要填充的，哪些是不需要填充的。

用户给出的PPT设计计划的格式说明：
模板类型: 
###
{template_type}
###

计划: 
###
{plan}
###
```

你要基于用户的计划，完善其中的细节，最终能够较好地填充到当前页面的模板元素中，用户将对应地进行替换，最终能够得到一个完整的PPT页面。

```
下面是一个输入的示例：
elements_info: [{{{{'type': 'TextBox', 'content': '单击此处输入你的项正文，文字是您思想的提炼，请尽量言简意赅的阐述观点，单击此处输入你的项正文'}}}},
{{{{'type': 'TextBox', 'content': '单击此处输入你的项正文，文字是您思想的提炼，请尽量言简意赅的阐述观点，单击此处输入你的项正文'}}}},
{{{{'type': 'TextBox', 'content': '单击此处输入你的项正文，文字是您思想的提炼，请尽量言简意赅的阐述观点，单击此处输入你的项正文'}}}},
{{{{'type': 'TextBox', 'content': '01'}}}},
{{{{'type': 'TextBox', 'content': '02'}}}},
{{{{'type': 'TextBox', 'content': '03'}}}},
{{{{'type': 'TextBox', 'content': '添加项标题'}}}},
{{{{'type': 'TextBox', 'content': '单击此处添加标题'}}}}]

模板类型: 文字内容页模板
计划: 需要填入引言的子标题内容，"机器学习技术分类"，同时围绕主题"机器学习技术分类"展开3项内容，满足顺序要求

```

你的执行思路应该是：
1. 识别当前页面的元素信息，找到对应的文本框
2. 识别用户的计划，将计划中的内容填充到文本框中
3. 重复1-2步骤，直到所有的文本框都填充完毕，并且满足了计划的要求

你要注意的是界面元素中有一些是不能修改的，比如说01、02、03等，这可能是固有的一些界面元素，你不能修改，而有一些你可以修改，你需要进行判断区分。
另外你要注意元素的含义，主要分为标题、小标题和内容，你要根据计划的要求，合理地填充内容，保证内容的完整性和连贯性，不要颠倒标题与内容！
 - 例如`单击此处输入你的项正文`是内容，而`单击此处添加标题`是标题，你要根据计划的要求，合理地填充内容，保证内容的完整性和连贯性。

对于文本框的数量你要进行判断，如果计划中的文本框数量大于实际的文本框数量，你需要修改计划，如果小于实际的文本框数量，你需要补充计划，最终你要确保每个文本框都被填充内容。
当然考虑到有些文本框可能填充的内容需要较少，而有一些则需要较多的内容，你需要根据计划的要求，合理地填充内容，保证内容的完整性和连贯性。

以下是输出字段说明：
element_idx: 元素的索引，对应于elements_info中的元素在列表中的index
content: 填充的内容

请输出一个json格式的结果，示例如下：
[{{{{
    "element_idx": 1,
    ”content": "有监督学习‌：XXX"
}}}}, {{{{
    "element_idx": 2,
    ”content": "无监督学习：XXX"
}}}}, {{{{
    "element_idx": 3,
    ”content": "半监督学习：XXX"
}}}}, {{{{
    "element_idx": 7,
    ”content": "机器学习技术"
}}}}, {{{{
    "element_idx": 8,
    ”content": "机器学习技术分类"
}}}}]
"""

    @classmethod
    def generate_ppt_with_basic_structure(cls, **kwargs):
        """
            基于用户简单输入所生成的基础PPT结构，生成PPT文件

        params:
            - template_id: 模板ID
        """
        from flask import g
        user_id = g.user['user_id']

        # 构造template
        template_id = kwargs.get('template_id')
        template = TemplateService.get(template_id=template_id)
        if template['code'] != RET.OK or template['totalCount'] != 1:
            return {'code': RET.DATAERR, 'message': '模板不存在', 'data': '模板不存在'}
        template = template['data'][0]
        template_path = os.path.join(cls.base_dir, template['location'])

        try:
            ppt_structure = ppt_structure_parser(template_path)
        except Exception as e:
            return {'code': RET.DATAERR, 'message': 'Template not found!', 'data': 'Template not found!'}

        ppt_generator = PPTGeneration(template_path)

        try:
            model = LLMService.get_model()
        except Exception as e:
            return {'code': RET.DATAERR, 'message': str(e), 'data': str(e)}

        format_ppt_element_enums = "\n".join([f"{e[0]}: {e[1]}" for e in PPT_ELEMENT_ENUMS.items()])

        try:
            # Step 1: Analysis the PPT structure and elements
            analysis_prompt = cls.get_prompt(
                prompt=cls.element_analysis_prompt.format(ppt_elements_enum=format_ppt_element_enums),
                user_message="{user_message}")

            analysis_chain = analysis_prompt | model
            res_elements = analysis_chain.invoke({"user_message": "请输出PPT模板结构：{}".format(ppt_structure)}).content
            print(res_elements)

            res_elements = fetch_json(res_elements)
            if not res_elements:
                return {'code': RET.DATAERR, 'message': '分析结果为空', 'data': '分析结果为空'}

            # Step 2: Plan the structure of the PPT generation

            # get structure
            ppt_id = kwargs.get('ppt_id')
            if ppt_id:
                ppt = PptBaseService.get(ppt_id=ppt_id)
                if ppt['code'] != RET.OK or ppt['totalCount'] != 1:
                    return {'code': RET.DATAERR, 'message': 'PPT不存在', 'data': 'PPT不存在'}
                ppt = ppt['data'][0]
                outline = ppt['outline']
            else:
                return {'code': RET.DATAERR, 'message': '大纲不存在，请重新对话', 'data': '大纲不存在，请重新对话'}

            plan_prompt = cls.get_prompt(
                prompt=cls.plan_ppt_prompt.format(ppt_elements_enum=format_ppt_element_enums),
                user_message="{user_message}")

            plan_chain = plan_prompt | model

            res_plan = plan_chain.invoke(
                {"user_message": "PPT的大纲：{outline}\nPPT模板元素：{elements}\n".format(outline=outline,
                                                                                        elements=res_elements)}).content

            print(res_plan)

            # Step 3: Generate each of the PPT file
            res_plan = fetch_json(res_plan)
            if not res_plan:
                return {'code': RET.DATAERR, 'message': '生成计划为空', 'data': '生成计划为空'}

            plans = json.loads(res_plan)
            for idx, p in enumerate(plans):
                element = ppt_structure[p['template_page'] - 1]
                formatted_element = ""
                for key, value in element.items():
                    if key == 'slide_number':
                        continue
                    elif key == 'elements':
                        formatted_element += "elements: ["
                        for i, e in enumerate(value):
                            formatted_element += "index: {index} | {type}: {content}\n".format(index=i, type=e['type'], content=e['content'])
                        formatted_element += "]\n"

                    else:
                        formatted_element += "{key}: {value}\n".format(key=key, value=value)

                # Document exists, then add the document information

                ppt_element = str(ppt_structure[p['template_page'] - 1]).replace('{', '{{').replace('}', '}}')
                template_type = PPT_ELEMENT_ENUMS[p['template_type']]
                plan = p['plan']
                if ppt['gen_type'] == 0:
                    gen_prompt = cls.get_prompt(
                        prompt=cls.single_ppt_generation.format(ppt_element=ppt_element,
                                                                template_type=template_type,
                                                                plan=plan),
                        user_message="{user_message}")
                    gen_chain = gen_prompt | model

                    res_generation = gen_chain.invoke({"user_message": "当前页面的模板元素: {element}\n模板类型:{template_type}\n计划: {plan}".format(
                            element=formatted_element, template_type=template_type, plan=plan)}).content
                # Document exists, then add the document information
                elif ppt['gen_type'] == 1:
                    vector_path = ppt['vector_path']
                    related = PptChatDocumentService.search_documents(vector_path, plan, limit=2)
                    related = "\n".join([r.page_content for r in related])
                    gen_prompt = cls.get_prompt(
                        prompt=cls.single_ppt_generation.format(ppt_element=ppt_element,
                                                                template_type=template_type,
                                                                plan=plan),
                        user_message="{user_message}")
                    gen_chain = gen_prompt | model

                    res_generation = gen_chain.invoke({"user_message":
                        "当前页面的模板元素: {element}\n模板类型:{template_type}\n计划: {plan}\n相关的文档材料：{document}\n在填充时结合文档材料与计划进行实施。".format(
                            element=formatted_element,
                            template_type=template_type,
                            plan=plan,
                            document=related)}).content
                else:
                    return {'code': RET.DATAERR, 'message': 'PPT类型错误', 'data': 'PPT类型错误'}



                print(res_generation)

                res_generation = fetch_json(res_generation)
                res_generation = json.loads(res_generation)

                # 生成PPT页
                try:
                    ppt_generator.generate_ppt(p['template_page'] - 1, res_generation)
                except Exception as e:
                    loggings.exception(1, e)

            # 保存PPT
            ppt_save_id = GenerateID.create_random_id()
            save_path = os.path.join(cls.ppt_save_dir, user_id, "{}.pptx".format(ppt_save_id))
            os.makedirs(os.path.join(cls.ppt_save_dir, user_id), exist_ok=True)

            ppt_generator.save_ppt(save_path)
        except Exception as e:
            loggings.exception(1, e)
            return {'code': RET.DATAERR, 'message': str(e), 'data': str(e)}

        # 写入数据库
        try:
            db.session.begin_nested()

            gen_id = GenerateID.create_random_id()
            ppt_gen_model = PptGen(
                gen_id=gen_id,
                ppt_id=ppt_id,
                template_id=template_id,
                location=os.path.join('uploads', 'generation', user_id, "{}.pptx".format(ppt_save_id)),
            )

            db.session.add(ppt_gen_model)

            ppt_base_model = db.session.query(PptBase).filter(*[ppt_id == ppt_id]).with_for_update()
            ppt_base_model.update({'ppt_location': ppt_gen_model.location, 'status': 1})

            db.session.commit()

            return {'code': RET.OK, 'message': '生成成功', 'data': {'gen_id': gen_id}}
        except Exception as e:
            db.session.rollback()
            loggings.exception(1, e)
            return {'code': RET.DBERR, 'message': str(e), 'data': str(e)}

    @classmethod
    def get_prompt(cls, prompt, user_message: str) -> ChatPromptTemplate:
        return ChatPromptTemplate.from_messages([("system", prompt), ("user", user_message)])

